from train.main import main


class Detector:
    def __init__(self):
        self.config = r'../train/default.yaml'

    # vgg16预测器
    def vgg16_predict(self, image_path):
        print("加载VGG16模型进行预测")
        vgg16_detector = main(self.config, "Vgg16", "predict",
                              ckpt=r"D:\pycharmProject\trash_classification\checkpoint\Vgg16\2\epoch95-loss0.07-acc0.83.pth")
        pred_str = vgg16_detector.predict(image_path)
        print(f"vgg16的预测结果为{pred_str}")
        return pred_str

    # resnet50预测器
    def resnet50_predict(self, image_path):
        print("加载resnet50模型进行预测")
        resnet50_detector = main(self.config, "ResNet50", "predict",
                                 ckpt=r"D:\pycharmProject\trash_classification\checkpoint\ResNet50\0\epoch95-loss0.07-acc0.84.pth")
        pred_str = resnet50_detector.predict(image_path)
        print(f"resnet50预测结果为{pred_str}")
        return pred_str
